/**
 * Copyright (C) 2010 - 2013 Harry Glasgow
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */
package com.googlecode.jaden.ui.dialog.manualtraining;

import com.googlecode.jaden.common.enums.LayerType;
import com.googlecode.jaden.common.other.Utils;
import com.googlecode.jaden.engine.Layer;
import com.googlecode.jaden.ui.UiUtils;
import static com.googlecode.jaden.ui.Constants.*;
import com.googlecode.jaden.ui.dialog.BaseDialog;

import javax.swing.*;
import javax.swing.event.ChangeEvent;
import javax.swing.event.ChangeListener;
import java.awt.event.ItemEvent;
import java.awt.event.ItemListener;
import java.util.logging.Logger;
import java.util.prefs.Preferences;
import java.util.Map;
import java.util.HashMap;
import java.util.concurrent.atomic.AtomicBoolean;

public class ConfigureManualTrainingDialog extends BaseDialog {

    private static final Logger logger = Logger.getLogger(ConfigureManualTrainingDialog.class.getName());
    public static final int MAX_LAYERS = 5;

    private final ManualTrainingDialog parent;
    private final int inputWidth;
    private final int outputWidth;

    private SpinnerNumberModel layersSpinnerModel;
    private DefaultComboBoxModel<String> layerComboModel;
    private JComboBox<String> layerCombo;
    private JComboBox<LayerType> layerTypeCombo;
    private JSpinner layerWidthSpinner;
    private SpinnerNumberModel layerWidthSpinnerModel;
    private SpinnerNumberModel learningRateSpinnerModel;
    private SpinnerNumberModel momentumSpinnerModel;
    private SpinnerNumberModel epochsSpinnerModel;

    private final Map<Integer, Layer> layers = new HashMap<Integer, Layer>();
    private AtomicBoolean updatingLayersCombo = new AtomicBoolean();
    private AtomicBoolean updatingLayerCombo = new AtomicBoolean();

    public ConfigureManualTrainingDialog(ManualTrainingDialog parent,
                                         Preferences preferences,
                                         Map<Integer, Layer> layers,
                                         double manualLearningRate,
                                         double manualMomentum,
                                         long manualEpochs,
                                         int inputWidth,
                                         int outputWidth) {
        super(parent, DIALOG_CONFIGURE_MANUAL_TRAINING_TITLE);
        this.parent = parent;
        this.inputWidth = inputWidth;
        this.outputWidth = outputWidth;
        this.layers.putAll(layers);

        layersSpinnerModel = new SpinnerNumberModel(1, 1, MAX_LAYERS , 1);
        JSpinner layersSpinner = new JSpinner(layersSpinnerModel);

        layerComboModel = new DefaultComboBoxModel<String>();
        layerCombo = new JComboBox<String>(layerComboModel);

        layerTypeCombo = new JComboBox<LayerType>();
        layerTypeCombo.addItem(LayerType.Soft);
        layerTypeCombo.addItem(LayerType.Normal);
        layerTypeCombo.addItem(LayerType.Sharp);

        layerWidthSpinnerModel = new SpinnerNumberModel(1, 1, 999, 1);
        layerWidthSpinner = new JSpinner(layerWidthSpinnerModel);

        learningRateSpinnerModel = new SpinnerNumberModel(0.1, 0.1, 1, 0.1);
        JSpinner learningRateSpinner = new JSpinner(learningRateSpinnerModel);

        momentumSpinnerModel = new SpinnerNumberModel(0.0, 0.0, 1.0, 0.1);
        JSpinner momentumSpinner = new JSpinner(momentumSpinnerModel);

        epochsSpinnerModel = new SpinnerNumberModel(1000, 1, 9999, 1);
        JSpinner epochsSpinner = new JSpinner(epochsSpinnerModel);

        JLabel layersLabel = new JLabel(Utils.translate(LABEL_LAYERS_TEXT));
        JLabel layerLabel = new JLabel(Utils.translate(LABEL_LAYER_TEXT));
        JLabel layerTypeLabel = new JLabel(Utils.translate(LABEL_LAYER_TYPE_TEXT));
        JLabel layerWidthLabel = new JLabel(Utils.translate(LABEL_LAYER_WIDTH_TEXT));
        JLabel learningRateLabel = new JLabel(Utils.translate(LABEL_LEARNING_RATE_TEXT));
        JLabel momentumLabel = new JLabel(Utils.translate(LABEL_MOMENTUM_TEXT));
        JLabel epochsLabel = new JLabel(Utils.translate(LABEL_EPOCHS_TEXT));

        UiUtils.enWidth(layersLabel, layerLabel, layerTypeLabel, layerWidthLabel,
                learningRateLabel, momentumLabel, epochsLabel);

        // Hack. The momentum spinner does not display enough width if its 
        // original value is zero. So force it to be the same size as
        // the learning rate spinner.
        UiUtils.enWidth(learningRateSpinner, momentumSpinner);

        JPanel mainPanel = UiUtils.verticalList(UiUtils.horizontalList(layersLabel, layersSpinner),
                new JSeparator(SwingConstants.HORIZONTAL),
                UiUtils.horizontalList(layerLabel, layerCombo),
                UiUtils.horizontalList(layerTypeLabel, layerTypeCombo),
                UiUtils.horizontalList(layerWidthLabel, layerWidthSpinner),
                new JSeparator(SwingConstants.HORIZONTAL),
                UiUtils.horizontalList(learningRateLabel, learningRateSpinner),
                UiUtils.horizontalList(momentumLabel, momentumSpinner),
                UiUtils.horizontalList(epochsLabel, epochsSpinner));
        setMainPanel(mainPanel, preferences);

        layersSpinnerModel.setValue(layers.size());
        layersSpinner.addChangeListener(new ChangeListener() {
            public void stateChanged(ChangeEvent e) {
                handleLayersComboChange();
            }
        });
        handleLayersComboChange();

        handleLayerComboChange();
        layerCombo.addItemListener(new ItemListener() {
            public void itemStateChanged(ItemEvent e) {
                handleLayerComboChange();
            }
        });

        layerTypeCombo.addItemListener(new ItemListener() {
            public void itemStateChanged(ItemEvent e) {
                updateLayerType();
            }
        });
        updateLayerType();

        layerWidthSpinner.addChangeListener(new ChangeListener() {
            public void stateChanged(ChangeEvent e) {
                updateLayerWidth();
            }
        });
        updateLayerWidth();

        learningRateSpinnerModel.setValue(manualLearningRate);
        momentumSpinnerModel.setValue(manualMomentum);
        epochsSpinnerModel.setValue(manualEpochs);
    }

    /**
     * Number of layers changing.
     */
    private void handleLayersComboChange() {

        // Disable other events while this is going on.
        updatingLayersCombo.set(true);

        // How many layers we got now?
        int newNumberOfLayers = layersSpinnerModel.getNumber().intValue();

        // Remove excess.
        while (layerComboModel.getSize() > newNumberOfLayers) {
            layerComboModel.removeElementAt(layerComboModel.getSize() - 1);
        }

        // Add up to the required number.
        while (layerComboModel.getSize() < newNumberOfLayers) {
            layerComboModel.addElement(Utils.translate(LABEL_LAYER_TEXT) +
                    (1 + layerComboModel.getSize()));
        }

        logger.info("input " + inputWidth);
        for (Integer integer : layers.keySet()) {
            Layer layer = layers.get(integer);
            logger.info(integer  + " was " + layer.getWeights()[0].length + " to " + layer.getLayerWidth());
        }
        logger.info("output " + outputWidth);

        // Delete excess layers.
        for (int i = 0; i < MAX_LAYERS; i++) {
            if (i >= newNumberOfLayers) {
                layers.remove(i);
            }
        }

        // Check width mappings.
        int previousWidth = inputWidth;
        for (int i = 0; i < newNumberOfLayers; i++) {
            Layer layer = layers.get(i);
            if (layer == null) {
                // Not there yet. Add it.
                layer = new Layer(LayerType.Normal, new double[previousWidth][previousWidth], new double[previousWidth]);
                layers.put(i, layer);
            } else if (layer.getWeights()[0].length != previousWidth) {
                // Misconfigured. Replace it.
                layer = new Layer(layer.getLayerType(), new double[layer.getLayerWidth()][previousWidth], new double[layer.getLayerWidth()]);
                layers.put(i, layer);
            }
            previousWidth = layer.getLayerWidth();
        }

        // Check last.
        int lastIndex = layers.size() - 1;
        Layer lastLayer = layers.get(lastIndex);
        if (lastLayer.getLayerWidth() != outputWidth) {
            Layer newLastLayer = new Layer(lastLayer.getLayerType(), new double[outputWidth][previousWidth], new double[outputWidth]);
            layers.put(lastIndex, newLastLayer);
        }

        logger.info("input " + inputWidth);
        for (Integer integer : layers.keySet()) {
            Layer layer = layers.get(integer);
            logger.info(integer  + " now " + layer.getWeights()[0].length + " to " + layer.getLayerWidth());
        }
        logger.info("output " + outputWidth);

        // Re-enable other events.
        updatingLayersCombo.set(false);

        handleLayerComboChange();
    }

    private void handleLayerComboChange() {
        if (!updatingLayersCombo.get()) {
            updatingLayerCombo.set(true);
            int layerIndex = layerCombo.getSelectedIndex();
            Layer layer = layers.get(layerIndex);

            if (layer.getLayerType() == LayerType.Soft) {
                layerTypeCombo.setSelectedIndex(0);
            } else if (layer.getLayerType() == LayerType.Normal) {
                layerTypeCombo.setSelectedIndex(1);
            } else if (layer.getLayerType() == LayerType.Sharp) {
                layerTypeCombo.setSelectedIndex(2);
            } else {
                Utils.handleException(new IllegalStateException("Unknown layer type " + layer.getLayerType()));
            }

            layerWidthSpinnerModel.setValue(layer.getLayerWidth());
            boolean last = layerIndex == layersSpinnerModel.getNumber().intValue() - 1;
            layerWidthSpinner.setEnabled(!last);
            updatingLayerCombo.set(false);
        }
    }
    
    private void updateLayerType() {
        if (!updatingLayersCombo.get() && !updatingLayerCombo.get()) {
            int layerIndex = layerCombo.getSelectedIndex();
            Layer layer = layers.get(layerIndex);
            int typeSelection = layerTypeCombo.getSelectedIndex();
            if (typeSelection == 0) {
                Layer newLayer = new Layer(LayerType.Soft, layer.getWeights(), layer.getBiases());
                layers.put(layerIndex, newLayer);
            } else if (typeSelection == 1) {
                Layer newLayer = new Layer(LayerType.Normal, layer.getWeights(), layer.getBiases());
                layers.put(layerIndex, newLayer);
            } else if (typeSelection == 2) {
                Layer newLayer = new Layer(LayerType.Sharp, layer.getWeights(), layer.getBiases());
                layers.put(layerIndex, newLayer);
            } else {
                Utils.handleException(new IllegalStateException("Unknown layer type selection " + typeSelection));
            }
        }
    }

    private void updateLayerWidth() {
        if (!updatingLayersCombo.get() && !updatingLayerCombo.get()) {
            int layerIndex = layerCombo.getSelectedIndex();

            Layer layer = layers.get(layerIndex);
            int width = layerWidthSpinnerModel.getNumber().intValue();
            Layer newLayer = new Layer(layer.getLayerType(), new double[width][layer.getWeights()[0].length],
                    new double[width]);
            layers.put(layerIndex, newLayer);

            // Don't adjust the following layer if this is the last one.
            if (layerIndex < layers.size() - 1) {
                int followingLayerIndex = layerCombo.getSelectedIndex() + 1;
                layer = layers.get(followingLayerIndex);
                newLayer = new Layer(layer.getLayerType(), new double[layer.getLayerWidth()][width],
                        new double[layer.getLayerWidth()]);
                layers.put(followingLayerIndex, newLayer);
            }
        }
    }

    protected void doOk() {
        parent.updateManualTraining(layers, learningRateSpinnerModel.getNumber().doubleValue(),
                momentumSpinnerModel.getNumber().doubleValue(), epochsSpinnerModel.getNumber().intValue());
        dispose();
    }
}